Preparing OPD Teacher by Merging LoRA weights in verl-omni

There’s already a tool export_fsdp_lora_adapter() in verl_omni/utils/fsdp_utils.py that supports this. It directly aggregates lora_* tensors from FSDP checkpoints, and outputs standard PEFT format (adapter_config.json + adapter_model.safetensors).

$ python3 -c "
from verl_omni.utils.fsdp_utils import export_fsdp_lora_adapter
print(export_fsdp_lora_adapter(
    'checkpoints/sd35_medium_ocr_distill/global_step_100/actor',
    output_dir='checkpoints/sd35_ocr_lora_adapter',
    base_model_name_or_path='stabilityai/stable-diffusion-3.5-medium', # or absolute path
))"

I did meet some problems here¹. The solution is that, we need to first edit the exported adapter_config.json in checkpoints/sd35_ocr_lora_adapter:

1. set task_type key to null
2. modify the term "0" of key target_modules to to_out.0

Then, we should merge the exported LoRA and the transformer part through PEFT.

1. Copy the stablilityai/stable-diffusion-3.5-medium folder to checkpoints/sd35_ocr_merged_teacher
2. Run the following Python script to load the transformer part with SD3Transformer2DModel and merge via PeftModel.from_pretrained then merge_and_unload.

import shutil, torch
from diffusers import SD3Transformer2DModel
from peft import PeftModel

base = "stabilityai/stable-diffusion-3.5-medium"
out = "checkpoints/sd35_ocr_merged_teacher"
lora = "checkpoints/sd35_ocr_lora_adapter"

tf = SD3Transformer2DModel.from_pretrained(
    base,
    subfolder="transformer",
    torch_dtype=torch.bfloat16
)
tf = PeftModel.from_pretrained(tf, lora)
              .merge_and_unload()

tf.save_pretrained(f"{out}/transformer")  # Only overwrite the transformer part

You may verify by running sha256sum ....../transformer/diffusion_pytorch_model.safetensors.

And we’re done! We can then pass TEACHER_PATH=/path/to/sd35_ocr_merged_teacher as argument to run Diffusion OPD with Stable Diffusion 3.5 Medium with single teacher.